import os
import requests
import time
from pathlib import Path
import json
import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageFont
import io
from config import *  # 导入所有配置

def verify_api_connection():
    """验证API连接是否正常"""
    try:
        url = f"{COMPREFACE_BASE_URL}/api/v1/recognition/subjects"
        headers = {
            'x-api-key': RECOGNITION_API_KEY
        }
        
        print(f"尝试连接: {url}")
        response = requests.get(url, headers=headers, timeout=TIMEOUT)
        
        print(f"API响应状态码: {response.status_code}")
        if response.status_code == 200:
            print("✓ API连接正常")
            return True
        else:
            print(f"✗ API连接异常，状态码: {response.status_code}")
            print(f"响应内容: {response.text}")
            return False
    except requests.exceptions.RequestException as e:
        print(f"✗ API连接失败: {str(e)}")
        return False

def detect_and_recognize_faces(image_path, retry=0):
    """检测并识别图片中的所有人脸"""
    # 第一步：检测人脸位置
    detection_url = f"{COMPREFACE_BASE_URL}/api/v1/detection/detect"
    detection_headers = {
        'x-api-key': DETECTION_API_KEY
    }
    
    # 第二步：识别人脸
    recognition_url = f"{COMPREFACE_BASE_URL}/api/v1/recognition/recognize"
    recognition_headers = {
        'x-api-key': RECOGNITION_API_KEY
    }
    
    try:
        with open(image_path, 'rb') as img_file:
            img_data = img_file.read()
            
        # 检测人脸
        print(f"正在检测图片中的人脸: {image_path}")
        detection_files = {'file': (os.path.basename(image_path), img_data, 'image/jpeg')}
        detection_data = {'det_prob_threshold': str(DET_PROB_THRESHOLD)}
        
        detection_response = requests.post(
            detection_url,
            headers=detection_headers,
            files=detection_files,
            data=detection_data,
            timeout=TIMEOUT
        )
        
        if detection_response.status_code != 200:
            print(f"✗ 人脸检测失败，状态码: {detection_response.status_code}")
            print(f"错误信息: {detection_response.text}")
            return None
            
        detection_result = detection_response.json()
        if 'result' not in detection_result or len(detection_result['result']) == 0:
            print("✗ 未检测到人脸")
            return None
            
        faces_count = len(detection_result['result'])
        print(f"✓ 检测到 {faces_count} 个人脸")
        
        # 识别人脸
        print("正在识别人脸...")
        recognition_files = {'file': (os.path.basename(image_path), img_data, 'image/jpeg')}
        recognition_data = {
            'limit': '5',  # 每个人脸返回最多5个可能的匹配
            'det_prob_threshold': str(DET_PROB_THRESHOLD),
            'prediction_count': '1'  # 每个人脸只返回最可能的1个匹配
        }
        
        recognition_response = requests.post(
            recognition_url,
            headers=recognition_headers,
            files=recognition_files,
            data=recognition_data,
            timeout=TIMEOUT
        )
        
        if recognition_response.status_code != 200:
            print(f"✗ 人脸识别失败，状态码: {recognition_response.status_code}")
            print(f"错误信息: {recognition_response.text}")
            return detection_result
            
        recognition_result = recognition_response.json()
        print(f"识别结果: {json.dumps(recognition_result, ensure_ascii=False, indent=2)[:300]}...")
        
        return {
            'detection': detection_result,
            'recognition': recognition_result
        }
                
    except requests.exceptions.Timeout:
        print(f"请求超时: {image_path}")
        if retry < MAX_RETRIES:
            print(f"尝试重试 ({retry+1}/{MAX_RETRIES})...")
            time.sleep(2)  # 等待2秒后重试
            return detect_and_recognize_faces(image_path, retry+1)
        return None
        
    except requests.exceptions.ConnectionError as e:
        print(f"连接错误: {str(e)}")
        print("请检查CompreFace服务是否正在运行，以及API地址是否正确")
        if retry < MAX_RETRIES:
            print(f"尝试重试 ({retry+1}/{MAX_RETRIES})...")
            time.sleep(2)  # 等待2秒后重试
            return detect_and_recognize_faces(image_path, retry+1)
        return None
        
    except Exception as e:
        print(f"发生错误: {str(e)}")
        return None

def save_results_to_txt(image_path, results):
    """将识别结果保存到txt文件中"""
    if not results:
        return None
        
    output_txt_path = f"{os.path.splitext(image_path)[0]}_result.txt"
    
    with open(output_txt_path, 'w', encoding='utf-8') as f:
        f.write(f"图片路径: {image_path}\n")
        f.write(f"处理时间: {time.strftime('%Y-%m-%d %H:%M:%S')}\n\n")
        
        # 写入检测结果
        if 'detection' in results:
            detection_result = results['detection']
            if 'result' in detection_result:
                faces_count = len(detection_result['result'])
                f.write(f"检测到 {faces_count} 个人脸\n\n")
                
                # 写入每个人脸的详细信息
                f.write("人脸详细信息:\n")
                for i, face in enumerate(detection_result['result']):
                    box = face['box']
                    f.write(f"人脸 #{i+1}:\n")
                    f.write(f"  位置: x={box['x_min']}~{box['x_max']}, y={box['y_min']}~{box['y_max']}\n")
                    f.write(f"  检测概率: {box.get('probability', 0):.4f}\n")
                    
                    # 添加识别结果（如果有）
                    if 'recognition' in results and 'result' in results['recognition']:
                        # 查找对应的识别结果
                        for rec_face in results['recognition']['result']:
                            rec_box = rec_face.get('box', {})
                            
                            # 检查是否是同一个人脸（通过比较框的位置）
                            if (abs(rec_box.get('x_min', 0) - box['x_min']) < 10 and 
                                abs(rec_box.get('y_min', 0) - box['y_min']) < 10):
                                
                                # 获取识别结果
                                if 'subjects' in rec_face and len(rec_face['subjects']) > 0:
                                    f.write("  识别结果:\n")
                                    for j, subject in enumerate(rec_face['subjects']):
                                        f.write(f"    候选 #{j+1}: {subject['subject']} (相似度: {subject['similarity']:.4f})\n")
                                else:
                                    f.write("  识别结果: 未能识别（库中可能没有匹配的人脸）\n")
                                break
                        else:
                            f.write("  识别结果: 未能识别（可能由于人脸角度、光线或遮挡等问题）\n")
                    f.write("\n")
        
        # 添加未识别人脸的可能原因
        if 'detection' in results and 'recognition' in results:
            det_count = len(results['detection'].get('result', []))
            rec_count = len(results['recognition'].get('result', []))
            recognized_count = sum(
                1 for face in results['recognition'].get('result', []) 
                if 'subjects' in face and len(face['subjects']) > 0
            )
            
            f.write("\n总结:\n")
            f.write(f"检测到的人脸数量: {det_count}\n")
            f.write(f"识别处理的人脸数量: {rec_count}\n")
            f.write(f"成功识别的人脸数量: {recognized_count}\n")
            
            if det_count > rec_count:
                f.write(f"\n注意: 有 {det_count - rec_count} 个检测到的人脸未被识别系统处理，可能原因:\n")
                f.write("1. 人脸太小或不清晰\n")
                f.write("2. 人脸角度不适合识别（侧脸、低头等）\n")
                f.write("3. 人脸被部分遮挡\n")
                f.write("4. 光线条件不佳\n")
            
            if rec_count > recognized_count:
                f.write(f"\n注意: 有 {rec_count - recognized_count} 个处理的人脸未能成功匹配到已知人物，可能原因:\n")
                f.write("1. 该人物未在系统中注册\n")
                f.write("2. 人脸特征与库中样本差异较大（年龄、发型、表情等变化）\n")
                f.write("3. 相似度低于系统阈值\n")
    
    print(f"识别结果已保存到: {output_txt_path}")
    return output_txt_path

def optimize_label_positions(faces, image_width, image_height):
    """优化标签位置，避免重叠"""
    # 按照y坐标排序人脸，从上到下
    sorted_faces = sorted(faces, key=lambda f: f['box']['y_min'])
    
    # 为每个人脸分配标签位置
    for i, face in enumerate(sorted_faces):
        box = face['box']
        face_width = box['x_max'] - box['x_min']
        face_height = box['y_max'] - box['y_min']
        
        # 默认位置（人脸上方）
        position = {
            'x': box['x_min'],
            'y': box['y_min'] - 25,
            'align': 'left',
            'vertical_offset': 0
        }
        
        # 检查是否靠近顶部边缘
        if box['y_min'] < 40:
            # 如果靠近顶部，将标签放在人脸下方
            position['y'] = box['y_max'] + 5
        
        # 检查是否靠近左侧边缘
        if box['x_min'] < 10:
            # 如果靠近左侧，将标签右对齐
            position['x'] = box['x_max']
            position['align'] = 'right'
        
        # 检查是否靠近右侧边缘
        if box['x_max'] > image_width - 10:
            # 如果靠近右侧，将标签左对齐
            position['x'] = box['x_min']
            position['align'] = 'left'
        
        # 处理可能的重叠
        for j in range(i):
            prev_face = sorted_faces[j]
            if 'label_position' not in prev_face:
                continue
                
            prev_pos = prev_face['label_position']
            
            # 简单检查是否重叠（垂直方向）
            if abs(position['y'] - prev_pos['y']) < 25:
                # 如果重叠，向下移动当前标签
                position['vertical_offset'] += 25
                position['y'] += position['vertical_offset']
        
        face['label_position'] = position
    
    return sorted_faces

def draw_faces_on_image(image_path, results):
    """在图片上绘制人脸框和识别结果，将所有识别结果统一放在图片底部"""
    if not results or 'detection' not in results:
        print("没有可用的检测结果，无法绘制")
        return None
        
    # 加载图片
    image = Image.open(image_path).convert("RGB")  # 确保图像是RGB模式
    image_width, image_height = image.size
    draw = ImageDraw.Draw(image)
    
    # 尝试加载中文字体
    try:
        # 在Windows系统上尝试加载中文字体
        font = ImageFont.truetype("simhei.ttf", 16)  # 使用黑体，字体稍小
        title_font = ImageFont.truetype("simhei.ttf", 20)  # 标题字体稍大
    except IOError:
        try:
            # 尝试其他常见中文字体
            font = ImageFont.truetype("simsun.ttc", 16)
            title_font = ImageFont.truetype("simsun.ttc", 20)
        except IOError:
            # 如果无法加载中文字体，使用默认字体
            font = ImageFont.load_default()
            title_font = ImageFont.load_default()
            print("警告：无法加载中文字体，将使用默认字体")
    
    # 获取检测和识别结果
    detection_result = results['detection']
    recognition_result = results.get('recognition', None)
    
    # 准备人脸数据
    faces = []
    for face in detection_result['result']:
        face_data = {'box': face['box']}
        
        # 如果有识别结果，添加标签
        label = "未知人物"
        similarity = 0
        
        if recognition_result and 'result' in recognition_result:
            # 查找对应的识别结果
            for rec_face in recognition_result['result']:
                rec_box = rec_face.get('box', {})
                
                # 检查是否是同一个人脸（通过比较框的位置）
                if (abs(rec_box.get('x_min', 0) - face['box']['x_min']) < 10 and 
                    abs(rec_box.get('y_min', 0) - face['box']['y_min']) < 10):
                    
                    # 获取识别结果
                    if 'subjects' in rec_face and len(rec_face['subjects']) > 0:
                        subject = rec_face['subjects'][0]
                        label = subject['subject']
                        similarity = subject['similarity']
                    break
        
        face_data['label'] = label
        face_data['similarity'] = similarity
        faces.append(face_data)
    
    # 给每个人脸分配一个数字标识
    for i, face in enumerate(faces):
        face['id'] = i + 1
    
    # 绘制人脸框和数字标识
    for face in faces:
        box = face['box']
        face_id = face['id']
        
        # 绘制人脸框
        draw.rectangle(
            [(box['x_min'], box['y_min']), (box['x_max'], box['y_max'])],
            outline=(0, 255, 0),
            width=2
        )
        
        # 绘制人脸ID标识（在人脸框左上角）
        id_text = str(face_id)
        id_bbox = draw.textbbox((0, 0), id_text, font=title_font)
        id_width = id_bbox[2] - id_bbox[0]
        id_height = id_bbox[3] - id_bbox[1]
        
        # 绘制ID背景
        draw.rectangle(
            [(box['x_min'] - 2, box['y_min'] - id_height - 4), 
             (box['x_min'] + id_width + 2, box['y_min'] - 2)],
            fill=(0, 255, 0)
        )
        
        # 绘制ID文本
        draw.text(
            (box['x_min'], box['y_min'] - id_height - 2),
            id_text,
            fill=(0, 0, 0),
            font=title_font
        )
    
    # 计算底部信息框的大小
    info_box_height = 0
    info_padding = 10
    line_height = 20
    
    # 标题行
    title_text = f"识别结果 (共 {len(faces)} 个人脸)"
    info_box_height += line_height + info_padding
    
    # 每个人脸的信息行
    for face in faces:
        # 人脸ID和标签
        label = face['label']
        if len(label) > 30:  # 如果标签太长，分成多行
            info_box_height += (len(label) // 30 + 1) * line_height
        else:
            info_box_height += line_height
            
        # 相似度行（如果有）
        if face['similarity'] > 0:
            info_box_height += line_height
            
        # 添加一点额外间距
        info_box_height += 5
    
    # 最后添加一些底部填充
    info_box_height += info_padding
    
    # 创建一个新的更大的图像，包含原图和底部信息框
    new_image = Image.new('RGB', (image_width, image_height + info_box_height), (255, 255, 255))
    new_image.paste(image, (0, 0))
    new_draw = ImageDraw.Draw(new_image)
    
    # 绘制底部信息框的背景
    new_draw.rectangle(
        [(0, image_height), (image_width, image_height + info_box_height)],
        fill=(240, 240, 240)
    )
    
    # 绘制标题
    new_draw.text(
        (info_padding, image_height + info_padding),
        title_text,
        fill=(0, 0, 0),
        font=title_font
    )
    
    # 绘制每个人脸的信息
    y_pos = image_height + info_padding + line_height + 5
    for face in faces:
        face_id = face['id']
        label = face['label']
        similarity = face['similarity']
        
        # 绘制人脸ID标识背景
        id_bbox = new_draw.textbbox((0, 0), str(face_id), font=title_font)
        id_width = id_bbox[2] - id_bbox[0]
        id_height = id_bbox[3] - id_bbox[1]
        
        new_draw.rectangle(
            [(info_padding, y_pos), (info_padding + id_width + 4, y_pos + id_height)],
            fill=(0, 255, 0)
        )
        
        # 绘制人脸ID
        new_draw.text(
            (info_padding + 2, y_pos),
            str(face_id),
            fill=(0, 0, 0),
            font=title_font
        )
        
        # 绘制标签
        if len(label) > 30:
            # 分行显示长标签
            for i in range(0, len(label), 30):
                segment = label[i:i+30]
                new_draw.text(
                    (info_padding + id_width + 15, y_pos),
                    segment,
                    fill=(0, 0, 0),
                    font=font
                )
                y_pos += line_height
        else:
            new_draw.text(
                (info_padding + id_width + 15, y_pos),
                label,
                fill=(0, 0, 0),
                font=font
            )
            y_pos += line_height
        
        # 绘制相似度（如果有）
        if similarity > 0:
            new_draw.text(
                (info_padding + id_width + 15, y_pos),
                f"相似度: {similarity:.4f}",
                fill=(0, 0, 150),
                font=font
            )
            y_pos += line_height
        
        # 添加一点额外间距
        y_pos += 5
    
    # 保存结果图片
    output_path = f"{os.path.splitext(image_path)[0]}_result{os.path.splitext(image_path)[1]}"
    new_image.save(output_path)
    print(f"已保存标注结果到: {output_path}")
    
    return output_path

def is_image_file(file_path):
    """检查文件是否为支持的图片格式"""
    return Path(file_path).suffix.lower() in IMAGE_EXTENSIONS

def process_images(image_paths):
    """处理指定的图片列表"""
    results = []
    
    for image_path in image_paths:
        if not os.path.exists(image_path):
            print(f"错误: 图片不存在 {image_path}")
            continue
            
        if not is_image_file(image_path):
            print(f"错误: 不支持的图片格式 {image_path}")
            continue
            
        print(f"\n处理图片: {image_path}")
        # 检测并识别人脸
        face_results = detect_and_recognize_faces(image_path)
        
        if face_results:
            # 保存结果到txt文件
            txt_path = save_results_to_txt(image_path, face_results)
            
            # 在图片上绘制结果
            output_path = draw_faces_on_image(image_path, face_results)
            if output_path:
                results.append({
                    'image_path': image_path,
                    'output_path': output_path,
                    'txt_path': txt_path,
                    'results': face_results
                })
    
    return results

def main():
    """主函数"""
    try:
        print(f"使用人脸检测API密钥: {DETECTION_API_KEY[:8]}...{DETECTION_API_KEY[-8:]}")
        print(f"使用人脸识别API密钥: {RECOGNITION_API_KEY[:8]}...{RECOGNITION_API_KEY[-8:]}")
        
        # 验证API连接
        print("正在验证API连接...")
        if not verify_api_connection():
            print("API连接验证失败，请检查CompreFace服务是否正在运行")
            print(f"当前API地址: {COMPREFACE_BASE_URL}")
            return
        
        # 处理图片
        if len(os.sys.argv) > 1:
            # 如果提供了命令行参数，使用参数作为图片路径
            image_paths = os.sys.argv[1:]
            print(f"从命令行参数获取图片路径: {image_paths}")
        else:
            # 否则使用face_images目录中的所有图片
            image_dir = "face_images"
            if not os.path.exists(image_dir):
                print(f"错误: 目录不存在 {image_dir}")
                return
                
            image_paths = [
                os.path.join(image_dir, f) 
                for f in os.listdir(image_dir) 
                if is_image_file(os.path.join(image_dir, f))
            ]
            print(f"从目录 {image_dir} 获取图片路径: {image_paths}")
        
        # 处理图片
        results = process_images(image_paths)
        
        # 输出结果摘要
        print("\n处理结果摘要:")
        for i, result in enumerate(results):
            print(f"{i+1}. 图片: {result['image_path']}")
            print(f"   输出图片: {result['output_path']}")
            print(f"   输出文本: {result['txt_path']}")
            
            if 'results' in result and 'recognition' in result['results']:
                det_result = result['results']['detection']
                rec_result = result['results']['recognition']
                
                det_count = len(det_result.get('result', []))
                rec_count = len(rec_result.get('result', []))
                recognized_count = sum(
                    1 for face in rec_result.get('result', []) 
                    if 'subjects' in face and len(face['subjects']) > 0
                )
                
                print(f"   检测到 {det_count} 个人脸，处理了 {rec_count} 个人脸，成功识别 {recognized_count} 个人脸")
                
                if det_count > rec_count:
                    print(f"   注意: 有 {det_count - rec_count} 个人脸未被识别系统处理")
        
        print("\n所有图片处理完成!")
        
    except Exception as e:
        print(f"发生错误: {str(e)}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    main()
